from sklearn.feature_extraction.text import TfidfVectorizer
from gq.trec_qa.data_processing import load_saved_data
from sklearn import tree
from sklearn.neighbors import KNeighborsClassifier
from sklearn.naive_bayes import MultinomialNB
from sklearn.linear_model import LogisticRegression
from sklearn import svm
from sklearn.metrics import accuracy_score
from sklearn.metrics import precision_score, recall_score, f1_score

x_train, y_train, x_test, y_test = load_saved_data()
tfidf_model = TfidfVectorizer(binary=False, token_pattern=r"(?u)\b\w+\b")
x_train = tfidf_model.fit_transform(x_train)
x_test = tfidf_model.transform(x_test)


def SVM():
    """
    SVM
    :return:
    """
    svm1 = svm.LinearSVC()
    svm1.fit(x_train, y_train)
    pred = svm1.predict(x_test)
    # pred_label = svm1.predict(text)
    precision, recall, f1 = computer_result(y_test, pred)
    # print(classification_report(y_test, pred, digits=4))
    return precision, recall, f1


def LR():
    """
    LR
    :return:
    """
    lr = LogisticRegression()
    lr.fit(x_train, y_train)
    pred = lr.predict(x_test)
    # pred_label = lr.predict(text)
    precision, recall, f1 = computer_result(y_test, pred)
    # print(classification_report(y_test, pred, digits=4))
    return precision, recall, f1


def KNN():
    """
    KNN
    :return:
    """
    knn = KNeighborsClassifier(n_neighbors=3)
    knn = knn.fit(x_train, y_train)
    pred = knn.predict(x_test)
    # pred_label = knn.predict(text)
    precision, recall, f1 = computer_result(y_test, pred)
    # print(classification_report(y_test, pred,digits=4))
    return precision, recall, f1


def Multi():
    """
    MultinomialNB
    :return:
    """
    classifier = MultinomialNB(alpha=0.001)
    classifier.fit(x_train, y_train)
    pred = classifier.predict(x_test)
    # pred_label = classifier.predict(text)
    precision, recall, f1 = computer_result(y_test, pred)
    # print(classification_report(y_test, pred, digits=4))
    return precision, recall, f1


def DecisionTree():
    """
    DecisionTree
    :return:
    """
    dtree = tree.DecisionTreeClassifier()
    dtree.fit(x_train, y_train)
    pred = dtree.predict(x_test)
    # pred_label = dtree.predict(text)
    # print(classification_report(y_test, pred, digits=4))
    precision, recall, f1 = computer_result(y_test, pred)
    return precision, recall, f1


def computer_result(y_true, y_pred):
    acc = accuracy_score(y_true, y_pred)
    precision = precision_score(y_true, y_pred, average='weighted')
    recall = recall_score(y_true, y_pred, average='weighted')
    f1 = f1_score(y_true, y_pred, average='weighted')
    return '%.3f' % precision, '%.3f' % recall, '%.3f' % f1


if __name__ == '__main__':
    print("svm", SVM())
    print('lr', LR())
    print('knn', KNN())
    print('multi', Multi())
    print('dtree', DecisionTree())
